{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ur8xi4C7S06n"
   },
   "outputs": [],
   "source": [
    "# Copyright 2025 Google LLC\n",
    "#\n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "#     https://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JAPoU8Sm5E6e"
   },
   "source": [
    "# MetaMath with Vertex AI Open Source Model Tuning\n",
    "\n",
    "<table align=\"left\">\n",
    "  <td style=\"text-align: center\">\n",
    "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/open-models/fine-tuning/get_started_with_oss_tuning_on_vertexai.ipynb\">\n",
    "      <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
    "    </a>\n",
    "  </td>\n",
    "  <td style=\"text-align: center\">\n",
    "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fopen-models%2Ffine-tuning%2Fget_started_with_oss_tuning_on_vertexai.ipynb\">\n",
    "      <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
    "    </a>\n",
    "  </td>\n",
    "  <td style=\"text-align: center\">\n",
    "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/open-models/fine-tuning/get_started_with_oss_tuning_on_vertexai.ipynb\">\n",
    "      <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\n",
    "    </a>\n",
    "  </td>\n",
    "  <td style=\"text-align: center\">\n",
    "    <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/open-models/fine-tuning/get_started_with_oss_tuning_on_vertexai.ipynb\">\n",
    "      <img width=\"32px\" src=\"https://raw.githubusercontent.com/primer/octicons/refs/heads/main/icons/mark-github-24.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
    "    </a>\n",
    "  </td>\n",
    "</table>\n",
    "\n",
    "<div style=\"clear: both;\"></div>\n",
    "\n",
    "<b>Share to:</b>\n",
    "\n",
    "<a href=\"https://www.linkedin.com/sharing/share-offsite/?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/open-models/fine-tuning/get_started_with_oss_tuning_on_vertexai.ipynb\" target=\"_blank\">\n",
    "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/8/81/LinkedIn_icon.svg\" alt=\"LinkedIn logo\">\n",
    "</a>\n",
    "\n",
    "<a href=\"https://bsky.app/intent/compose?text=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/open-models/fine-tuning/get_started_with_oss_tuning_on_vertexai.ipynb\" target=\"_blank\">\n",
    "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/7/7a/Bluesky_Logo.svg\" alt=\"Bluesky logo\">\n",
    "</a>\n",
    "\n",
    "<a href=\"https://twitter.com/intent/tweet?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/open-models/fine-tuning/get_started_with_oss_tuning_on_vertexai.ipynb\" target=\"_blank\">\n",
    "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/5a/X_icon_2.svg\" alt=\"X logo\">\n",
    "</a>\n",
    "\n",
    "<a href=\"https://reddit.com/submit?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/open-models/fine-tuning/get_started_with_oss_tuning_on_vertexai.ipynb\" target=\"_blank\">\n",
    "  <img width=\"20px\" src=\"https://redditinc.com/hubfs/Reddit%20Inc/Brand/Reddit_Logo.png\" alt=\"Reddit logo\">\n",
    "</a>\n",
    "\n",
    "<a href=\"https://www.facebook.com/sharer/sharer.php?u=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/open-models/fine-tuning/get_started_with_oss_tuning_on_vertexai.ipynb\" target=\"_blank\">\n",
    "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\" alt=\"Facebook logo\">\n",
    "</a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "84f0f73a0f76"
   },
   "source": [
    "| Author(s) |\n",
    "| --- |\n",
    "| [Ivan Nardini](https://github.com/inardini) |"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tvgnzT1CKxrO"
   },
   "source": [
    "## Overview\n",
    "\n",
    "This notebook demonstrates how to reproduce the core ideas of the **MetaMath** paper by fine-tuning a Llama model on the `MetaMathQA` dataset using Vertex AI's managed service for open-source models.\n",
    "\n",
    "### Objective\n",
    "\n",
    "The goal is to leverage the `MetaMathQA` dataset to enhance the mathematical reasoning capabilities of a base Llama model. Specifically, this notebook focuses on reproducing the results for the **7B model** variant discussed in the paper, using the comparable **Llama 3.1 8B model** available on Vertex AI.\n",
    "\n",
    "We will cover the following steps:\n",
    "\n",
    "1.  **Prepare the Dataset**: Download the `MetaMathQA` dataset and convert it to the required JSON Lines (JSONL) format for Vertex AI.\n",
    "2.  **Fine-Tune the Model**: Configure and launch a managed fine-tuning job on Vertex AI using a Llama 3.1 8B model.\n",
    "3.  **Deploy & Evaluate**: Deploy the newly tuned model to a Vertex AI Endpoint and test its mathematical reasoning.\n",
    "4.  **Compare (Optional)**: Compare our model's output with the official pre-trained MetaMath model from Hugging Face.\n",
    "5.  **Run Official Evaluation (Advanced)**: Download our tuned model and run the official evaluation scripts from the MetaMath repository.\n",
    "\n",
    "\n",
    "### Citation\n",
    "\n",
    "```\n",
    "@article{yu2023metamath,\n",
    "  title={MetaMath: Bootstrap Your Own Mathematical Questions for Large Language Models},\n",
    "  author={Yu, Longhui and Jiang, Weisen and Shi, Han and Yu, Jincheng and Liu, Zhengying and Zhang, Yu and Kwok, James T and Li, Zhenguo and Weller, Adrian and Liu, Weiyang},\n",
    "  journal={arXiv preprint arXiv:2309.12284},\n",
    "  year={2023}\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "61RBz8LLbxCR"
   },
   "source": [
    "## Get started\n",
    "\n",
    "This initial section handles all the necessary setup, including authenticating your account, installing libraries, and configuring your Google Cloud environment."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sOwtBan9NObx"
   },
   "source": [
    "### (Optional) Choose your runtime.\n",
    "\n",
    "This tutorial has been designed to compare your tuned version with the MetaMath model by running some tests and the evaluation benchmarks.\n",
    "\n",
    "To run **ONLY** those optional sections of this tutorial, you need to have access more resources compared to the one are available in free Colab. Consider to create a Workbench instance on Vertex AI using `a2-highgpu-2g (Accelerator Optimized: 2 NVIDIA Tesla A100 GPUs, 24 vCPUs, 170GB RAM)`.\n",
    "\n",
    "Optionally, you can install `jupyterlab-nvdashboard` to visualize GPU usage metrics within your notebook environment."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "No17Cw5hgx12"
   },
   "source": [
    "### Install Google Gen AI SDK and other required packages\n",
    "\n",
    "Install the Python libraries needed for this tutorial.\n",
    "\n",
    "  * `google-cloud-aiplatform`: The official SDK for interacting with Vertex AI services like model tuning and deployment.\n",
    "  * `datasets`: A library from Hugging Face that makes it easy to download and manipulate datasets.\n",
    "  * `transformers`: A Hugging Face library used for downloading and running the official pre-trained MetaMath model for comparison.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tFy3H3aPgx12"
   },
   "outputs": [],
   "source": [
    "%pip install --upgrade --quiet --force-reinstall google-cloud-aiplatform>=1.105.0\n",
    "%pip install --upgrade --quiet --force-reinstall datasets transformers torch sentencepiece accelerate bitsandbytes vllm hf_transfer fraction tqdm numpy fire openai scipy jsonlines pandas pydantic crcmod"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dmWOrTJ3gx13"
   },
   "source": [
    "### Authenticate your notebook environment (Colab only)\n",
    "\n",
    "If you're running this notebook on Google Colab, run the cell below to authenticate your environment.\n",
    "\n",
    "This gives the notebook permission to access your Google Cloud resources on your behalf.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "NyKGtVQjgx13"
   },
   "outputs": [],
   "source": [
    "# import sys\n",
    "\n",
    "# if \"google.colab\" in sys.modules:\n",
    "#     from google.colab import auth\n",
    "\n",
    "#     auth.authenticate_user()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DF4l8DTdWgPY"
   },
   "source": [
    "### Set Google Cloud project information\n",
    "\n",
    "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
    "\n",
    "Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment).\n",
    "\n",
    "Then, you define essential configuration variables. All resources you create (like models and storage buckets) will be associated with your specified Google Cloud project and region.\n",
    "\n",
    "  * `PROJECT_ID`: Your unique Google Cloud project identifier.\n",
    "  * `REGION`: The geographic location where your resources will be created (e.g., `us-central1`).\n",
    "  * `BUCKET_URI`: A unique Google Cloud Storage (GCS) bucket. This will be our central location for storing the training dataset and the resulting model artifacts.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Nqwi-5ufWp_B"
   },
   "outputs": [],
   "source": [
    "# Use the environment variable if the user doesn't provide Project ID.\n",
    "import os\n",
    "\n",
    "import vertexai\n",
    "\n",
    "PROJECT_ID = \"[your-project-id]\"  # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n",
    "\n",
    "if not PROJECT_ID or PROJECT_ID == \"[your-project-id]\":\n",
    "    PROJECT_ID = str(os.environ.get(\"GOOGLE_CLOUD_PROJECT\"))\n",
    "\n",
    "LOCATION = \"us-central1\"\n",
    "\n",
    "BUCKET_NAME = \"[your-bucket-name]\"  # @param {type: \"string\", placeholder: \"[your-bucket-name]\", isTemplate: true}\n",
    "BUCKET_URI = f\"gs://{BUCKET_NAME}\"\n",
    "\n",
    "# Create the GCS bucket if it doesn't exist\n",
    "! gsutil mb -p {PROJECT_ID} -l {LOCATION} {BUCKET_URI}\n",
    "\n",
    "# Initialize the Vertex AI SDK. This authenticates our session and sets the default project and location.\n",
    "vertexai.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5303c05f7aa6"
   },
   "source": [
    "### Import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6fc324893334"
   },
   "outputs": [],
   "source": [
    "import gc\n",
    "import json\n",
    "import os\n",
    "import time\n",
    "import uuid\n",
    "\n",
    "import torch\n",
    "import vertexai\n",
    "from datasets import load_dataset\n",
    "from google.cloud import aiplatform\n",
    "from google.cloud.aiplatform_v1beta1.types import JobState\n",
    "from pydantic import BaseModel, Field\n",
    "from transformers import pipeline\n",
    "from vertexai.preview import model_garden\n",
    "from vertexai.preview.tuning import sft\n",
    "from vertexai.preview.tuning._tuning import SourceModel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kh9YeyTMXSA-"
   },
   "source": [
    "### Helpers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "AfWTScHyXTo3"
   },
   "outputs": [],
   "source": [
    "def save_to_jsonl(dataset, output_path):\n",
    "    \"\"\"Save dataset in JSONL format required by Vertex AI.\"\"\"\n",
    "    with open(output_path, \"w\") as f:\n",
    "        for example in dataset:\n",
    "            json.dump(example, f)\n",
    "            f.write(\"\\n\")\n",
    "    print(f\"Saved {len(dataset)} examples to {output_path}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KPMJku26fvcR"
   },
   "source": [
    "## Prepare the MetaMathQA Dataset\n",
    "\n",
    "The core of the MetaMath paper is its unique dataset, `MetaMathQA`. Here, we'll download it, process it into the required format, and upload it to our cloud storage bucket so the Vertex AI tuning service can access it.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "76OvI4emfyLI"
   },
   "source": [
    "### Download and Format the Dataset\n",
    "\n",
    "The Vertex AI tuning service expects the training data to be in a **JSON Lines (JSONL)** format, where each line is a separate JSON object. For chat models, each object should contain a `\"messages\"` field with a list of conversation turns. We'll download the dataset from Hugging Face and map its `question` and `response` columns to this required format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "bJOxPVu7gAC7"
   },
   "outputs": [],
   "source": [
    "# Load the MetaMathQA dataset from the Hugging Face Hub.\n",
    "# We'll use the 'GSM8K' configuration, which is a key part of the paper's contribution.\n",
    "dataset = load_dataset(\"meta-math/MetaMathQA\")[\"train\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Jr2nwdu0uQwh"
   },
   "outputs": [],
   "source": [
    "# Use the .train_test_split() method to create an 80/20 split.\n",
    "# 80% of the data will be for training, 20% for validation.\n",
    "# This easily satisfies the <25% requirement for the validation set.\n",
    "split_dataset = dataset.train_test_split(test_size=0.2, seed=42)\n",
    "\n",
    "# The result is a dictionary containing the two new splits.\n",
    "train_split = split_dataset[\"train\"]\n",
    "validation_split = split_dataset[\"test\"]\n",
    "\n",
    "# Limit validation dataset to less than 5000 rows for Vertex AI requirement\n",
    "if len(validation_split) > 5000:\n",
    "    validation_split = validation_split.shuffle(\n",
    "        seed=42\n",
    "    )  # Use a seed for reproducibility\n",
    "    validation_split = validation_split.select(range(4999))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "1T0m86eqf9o6"
   },
   "outputs": [],
   "source": [
    "# MetaMath's instruction template\n",
    "METAMATH_TEMPLATE = \"\"\"Below is an instruction that describes a task. Write a response that appropriately completes the request.\n",
    "\n",
    "### Instruction:\n",
    "{instruction}\n",
    "\n",
    "### Response:\"\"\"\n",
    "\n",
    "# Define a function to transform each example into the {\"prompt\": ..., \"completion\": ...} structure.\n",
    "def format_for_tuning(example):\n",
    "\n",
    "    query = example[\"query\"]\n",
    "    response = example[\"response\"]\n",
    "\n",
    "    instruction = METAMATH_TEMPLATE.format(instruction=query)\n",
    "\n",
    "    # Important: Add space before response for proper tokenization\n",
    "    return {\n",
    "        \"messages\": [\n",
    "            {\"role\": \"user\", \"content\": instruction},\n",
    "            {\"role\": \"assistant\", \"content\": f\" {response}\"},\n",
    "        ]\n",
    "    }\n",
    "\n",
    "\n",
    "# Apply the formatting function to the entire dataset.\n",
    "train_formatted_dataset = train_split.map(\n",
    "    format_for_tuning, remove_columns=train_split.column_names, num_proc=os.cpu_count()\n",
    ")\n",
    "validation_formatted_dataset = validation_split.map(\n",
    "    format_for_tuning,\n",
    "    remove_columns=validation_split.column_names,\n",
    "    num_proc=os.cpu_count(),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6JFCFh3As8Ka"
   },
   "outputs": [],
   "source": [
    "train_file_path = \"metamath_gsm8k_train.jsonl\"\n",
    "validation_file_path = \"metamath_gsm8k_validation.jsonl\"\n",
    "\n",
    "# Write the formatted training data to a local JSONL file.\n",
    "save_to_jsonl(train_formatted_dataset, train_file_path)\n",
    "save_to_jsonl(validation_formatted_dataset, validation_file_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BI3_W51PmqxU"
   },
   "source": [
    "### Upload Dataset to GCS\n",
    "\n",
    "The Vertex AI tuning service runs on Google's infrastructure and cannot directly access files in this notebook's local environment. Therefore, we must upload our formatted JSONL file to our GCS bucket."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "DpzE-ksCm_PE"
   },
   "outputs": [],
   "source": [
    "# Define the destination path in your GCS bucket.\n",
    "train_file_uri = f\"{BUCKET_URI}/datasets/metamath_gsm8k_train.jsonl\"\n",
    "validation_file_uri = f\"{BUCKET_URI}/datasets/metamath_gsm8k_validation.jsonl\"\n",
    "\n",
    "# Use the gsutil command-line tool to copy the local file to GCS.\n",
    "! gsutil -o GSUtil:parallel_composite_upload_threshold=150M -m cp {train_file_path} {train_file_uri}\n",
    "! gsutil -o GSUtil:parallel_composite_upload_threshold=150M -m cp {validation_file_path} {validation_file_uri}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Qzk4nMOJpy6q"
   },
   "source": [
    "## Configure and Launch the Fine-Tuning Job\n",
    "\n",
    "We'll define all the parameters according to the official MetaMath README for our fine-tuning job and then launch it.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0ut8XGrqrTRY"
   },
   "source": [
    "### Define configuration\n",
    "\n",
    "Set tuning parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "kVnrHDAhpxH6"
   },
   "outputs": [],
   "source": [
    "# This class groups all hyperparameters and provides documentation and default values.\n",
    "class MetaMathTuningConfig(BaseModel):\n",
    "    \"\"\"Configuration settings for the MetaMath fine-tuning job.\"\"\"\n",
    "\n",
    "    base_model: str = Field(\n",
    "        default=\"meta/llama3_1@llama-3.1-8b\",\n",
    "        description=\"The base model to fine-tune, corresponding to the 7B model in the paper.\",\n",
    "    )\n",
    "    tuning_mode: str = Field(\n",
    "        default=\"FULL\",\n",
    "        description=\"The tuning mode. We use 'FULL' to replicate the paper's method for the 7B model.\",\n",
    "    )\n",
    "    epochs: int = Field(\n",
    "        default=3,\n",
    "        description=\"Number of training epochs, as specified in the MetaMath paper.\",\n",
    "    )\n",
    "    learning_rate: float = Field(\n",
    "        default=2e-5,\n",
    "        description=\"The learning rate for the optimizer, matching the paper's value for full fine-tuning.\",\n",
    "    )\n",
    "\n",
    "\n",
    "# Create an instance of the configuration class.\n",
    "config = MetaMathTuningConfig()\n",
    "\n",
    "# Dynamically create paths that depend on runtime variables.\n",
    "output_uri = f\"{BUCKET_URI}/tuning-output/{uuid.uuid4()}\"\n",
    "model_artifacts_gcs_uri = os.path.join(\n",
    "    output_uri, \"postprocess/node-0/checkpoints/final\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_-EdCljErt7x"
   },
   "source": [
    "### Launch the full fine-tuning job\n",
    "\n",
    "This function sends our configuration to the Vertex AI service, which will provision machines and run the training.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "FWBih5jfrNaA"
   },
   "outputs": [],
   "source": [
    "source_model = SourceModel(base_model=config.base_model)\n",
    "\n",
    "sft_tuning_job = sft.preview_train(\n",
    "    source_model=source_model,\n",
    "    tuning_mode=config.tuning_mode,\n",
    "    epochs=config.epochs,\n",
    "    learning_rate=config.learning_rate,\n",
    "    train_dataset=train_file_uri,\n",
    "    validation_dataset=validation_file_uri,\n",
    "    output_uri=output_uri,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EKoGUHLNsTZs"
   },
   "source": [
    "### Monitor the Job\n",
    "\n",
    "The tuning job runs remotely on Google Cloud. The following code provides a convenient way to check the job's status from within the notebook without having to manually refresh the console. It will print an update every 10 minutes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "fMwlgrUtsc-C"
   },
   "outputs": [],
   "source": [
    "print(\n",
    "    \"Monitoring job... This will take several hours. You can safely close this notebook and come back later.\"\n",
    ")\n",
    "\n",
    "while not sft_tuning_job.state in [\n",
    "    JobState.JOB_STATE_CANCELLED,\n",
    "    JobState.JOB_STATE_FAILED,\n",
    "    JobState.JOB_STATE_SUCCEEDED,\n",
    "]:\n",
    "    time.sleep(600)  # Check status every 10 minutes\n",
    "    sft_tuning_job.refresh()\n",
    "    print(f\"Current job state: {str(sft_tuning_job.state.name)}\")\n",
    "\n",
    "print(f\"Job finished with state: {sft_tuning_job.state.name}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CIzIsfVMLN5k"
   },
   "source": [
    "## Deploy the Tuned Model\n",
    "\n",
    "Once the tuning job is complete, the new model \"lives\" as a set of files in your GCS bucket. To use it for inference, we must **deploy** it. This process loads the model onto a server with a GPU and creates a unique API endpoint that we can send prediction requests to."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "V9dym9GbLTNb"
   },
   "outputs": [],
   "source": [
    "# Define the hardware for our deployment. An L4 GPU is a cost-effective choice for a model of this size.\n",
    "machine_type = \"g2-standard-12\"\n",
    "accelerator_type = \"NVIDIA_L4\"\n",
    "accelerator_count = 1\n",
    "\n",
    "# Create a CustomModel object that points to our tuned model artifacts in GCS.\n",
    "tuned_model = model_garden.CustomModel(gcs_uri=model_artifacts_gcs_uri)\n",
    "\n",
    "# Deploy the model. This step provisions the hardware and can take 15-30 minutes.\n",
    "endpoint = tuned_model.deploy(\n",
    "    machine_type=machine_type,\n",
    "    accelerator_type=accelerator_type,\n",
    "    accelerator_count=accelerator_count,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uMhu2iwPVBpO"
   },
   "source": [
    "## Evaluate and Compare\n",
    "\n",
    "Now, the exciting part! Let's test our newly tuned model and compare its performance to the official version."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZKk-YQBkVGdO"
   },
   "source": [
    "### Evaluate Tuned Model\n",
    "\n",
    "We'll send a sample math problem to our deployed endpoint. The prompt format is critical; we use the exact template specified in the official MetaMath repository to ensure the model responds as expected."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "fl9PnNwFU0mS"
   },
   "outputs": [],
   "source": [
    "# The inference prompt for MetaMath models.\n",
    "prompt_template = \"Below is an instruction that describes a task. Write a response that appropriately completes the request.\\n\\n### Instruction:\\n{instruction}\\n\\n### Response: Let's think step by step.\"\n",
    "instruction = \"James buys 5 packs of beef that are 4 pounds each. The price of beef is $5.50 per pound. How much did he pay?\"\n",
    "\n",
    "# Define the prediction request payload.\n",
    "# We set a low temperature for more factual, less \"creative\" output.\n",
    "instances = [\n",
    "    {\n",
    "        \"prompt\": prompt_template.format(instruction=instruction),\n",
    "        \"max_tokens\": 250,\n",
    "        \"temperature\": 0.2,\n",
    "        \"top_p\": 1.0,\n",
    "        \"top_k\": 1,\n",
    "        \"raw_response\": True,\n",
    "    }\n",
    "]\n",
    "\n",
    "# Send the request to our endpoint.\n",
    "response = endpoint.predict(instances=instances, use_dedicated_endpoint=True)\n",
    "\n",
    "print(\"Response from tuned model\")\n",
    "for prediction in response.predictions:\n",
    "    print(prediction)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JHRtxQd3bq06"
   },
   "source": [
    "### (Optional) Compare with the official MetaMath Model\n",
    "\n",
    "\n",
    "To see how well our reproduction worked, we can compare its output to the official `MetaMath-7B-V1.0` model released by the authors on Hugging Face. This provides a valuable benchmark.\n",
    "\n",
    "**Note**: This step runs a large model locally on the notebook's machine and may require significant RAM, GPU and time to download and generate predictions.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PCTnUS6_b71g"
   },
   "outputs": [],
   "source": [
    "# Enable hf_transfer for parallel downloads\n",
    "os.environ[\"HF_HUB_ENABLE_HF_TRANSFER\"] = \"1\"\n",
    "\n",
    "# Load the official MetaMath 7B model from Hugging Face.\n",
    "official_pipe = pipeline(\n",
    "    \"text-generation\", model=\"meta-math/MetaMath-7B-V1.0\", device_map=\"auto\"\n",
    ")\n",
    "\n",
    "# Use the same prompt and instruction for a fair comparison.\n",
    "official_response = official_pipe(\n",
    "    prompt_template.format(instruction=instruction), max_new_tokens=250, do_sample=False\n",
    ")\n",
    "\n",
    "print(\"Response from the MetaMath Model\")\n",
    "print(official_response[0][\"generated_text\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dI2n45VTpZu_"
   },
   "source": [
    "### Clean runtime\n",
    "\n",
    "To clear GPU memory we explicitly delete the pipeline, free up unused memory from PyTorch cache and any memory occupied by objects that are no longer referenced."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "bWJnyKyhpZu_"
   },
   "outputs": [],
   "source": [
    "del official_pipe\n",
    "torch.cuda.empty_cache()\n",
    "gc.collect()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "t7PHkLLUVH4C"
   },
   "source": [
    "### (Optional) Run the Official Evaluation Scripts\n",
    "\n",
    "The single-prompt tests above are good for a qualitative check. To get the official `pass@1` benchmark scores reported in the paper, you must run the evaluation scripts from the MetaMath GitHub repository against the full test dataset. This is how academic results are formally measured."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "GF5RjQDqFGpK"
   },
   "source": [
    "#### Clone the MetaMath Repository\n",
    "This command downloads the evaluation scripts (`eval_gsm8k.py`) and the test data files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "C6Ix2T0WFbrg"
   },
   "outputs": [],
   "source": [
    "!git clone https://github.com/meta-math/MetaMath.git"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "VgCS_8peFNSc"
   },
   "source": [
    "#### Download Your Tuned Model from GCS\n",
    "The evaluation script runs locally, so it needs the model files on the notebook's machine. We'll copy them from our GCS output directory.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "UQGrW2C8Ffzx"
   },
   "outputs": [],
   "source": [
    "# Create a local directory to store the model.\n",
    "LOCAL_MODEL_PATH = \"./my_tuned_metamath_model\"\n",
    "!mkdir -p {LOCAL_MODEL_PATH}\n",
    "\n",
    "# Copy the model files from GCS to the local path. This can take several minutes.\n",
    "!gsutil -m cp -r {model_artifacts_gcs_uri}/* {LOCAL_MODEL_PATH}/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ky7SwphaFvQ1"
   },
   "source": [
    "#### Run the GSM8K Evaluation\n",
    "Finally, execute the official evaluation script, pointing it to your locally downloaded model. This will run the model on every question in the GSM8K test set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "fIG1Ek1pF0-_"
   },
   "outputs": [],
   "source": [
    "!python MetaMath/eval_gsm8k.py \\\n",
    "    --model {LOCAL_MODEL_PATH} \\\n",
    "    --data_file ./MetaMath/data/test/GSM8K_test.jsonl \\\n",
    "    --tensor_parallel_size 2 \\\n",
    "    --batch_size 32"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gXQ-q0cRF2LS"
   },
   "source": [
    "After running, the script will output the final `pass@1` accuracy score. You can compare this number directly to the results table in the MetaMath paper to see how well your model performed! Also if you are interested, you can run the same process to measure `math` metrics."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2a4e033321ad"
   },
   "source": [
    "## Cleaning up\n",
    "\n",
    "To avoid incurring ongoing charges for the deployed model and stored data, you must undeploy the endpoint and delete your GCS artifacts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "f3PTswIxGLy3"
   },
   "outputs": [],
   "source": [
    "delete_experiments = True\n",
    "delete_endpoint = True\n",
    "delete_bucket = True\n",
    "\n",
    "# Deleting experiment\n",
    "if delete_experiments:\n",
    "    experiment = aiplatform.Experiment.list()[0]\n",
    "    experiment.delete()\n",
    "\n",
    "# Deleting the endpoint itself removes the resource configuration.\n",
    "if delete_endpoint:\n",
    "    endpoint = aiplatform.Endpoint.list()[0]\n",
    "    endpoint.delete(force=True)\n",
    "\n",
    "# To fully clean up, you should also delete the model artifacts and dataset from your GCS bucket.\n",
    "# You can do this via the command line or the Google Cloud Console.\n",
    "if delete_bucket:\n",
    "    !gsutil -m rm -r {BUCKET_URI}"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "name": "get_started_with_oss_tuning_on_vertexai.ipynb",
   "toc_visible": true
  },
  "environment": {
   "kernel": "conda-base-py",
   "name": "workbench-notebooks.m131",
   "type": "gcloud",
   "uri": "us-docker.pkg.dev/deeplearning-platform-release/gcr.io/workbench-notebooks:m131"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "conda-base-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
